import torch
from model import AlexNet
from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt
import json
import os

# 预处理
data_transform = transforms.Compose(
    [transforms.Resize((224, 224)),
     transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

# 读取 class_indices.json
try:
    with open('./class_indices.json', 'r') as json_file:
        class_indict = json.load(json_file)
except Exception as e:
    print(e)
    exit(-1)

# 创建模型
model = AlexNet(num_classes=10)
# 加载模型权重
model_weight_path = "./AlexNet.pth"
model.load_state_dict(torch.load(model_weight_path, map_location=torch.device('cpu'), weights_only=True))

# 关闭 Dropout
model.eval()

# 测试数据集文件夹路径
test_dir = 'C:\\pycharmproject\\project1\\SVCM\\car_data\\test'  # 替换为你的测试数据集文件夹路径

# 遍历测试文件夹中的所有图像
for filename in os.listdir(test_dir):
    if filename.lower().endswith((".jpg", ".jpeg", ".png")):  # 检查文件扩展名
        image_path = os.path.join(test_dir, filename)
        img = Image.open(image_path)
        plt.imshow(img)
        plt.title(f"Predicting: {filename}")  # 显示图片名称
        # [N, C, H, W]
        img = data_transform(img)
        # expand batch dimension
        img = torch.unsqueeze(img, dim=0)

        with torch.no_grad():
            # predict class
            output = torch.squeeze(model(img))  # 将输出压缩，即压缩掉 batch 这个维度
            predict = torch.softmax(output, dim=0)
            predict_cla = torch.argmax(predict).numpy()
            print(f"Image: {filename}, Predicted class: {class_indict[str(predict_cla)]}, Probability: {predict[predict_cla].item():.4f}")
        #plt.show()